#include <iostream>
#include <sys/socket.h>
#include <sys/uio.h> // writev 需要的头文件
#include <unistd.h>  // close, write
#include <string.h>  // strlen
#include <fcntl.h>   // fcntl
#include <errno.h>

// 设置文件描述符为非阻塞
int setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) return -1;
    return fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

int main() {
    // 创建一对 socket 用于模拟客户端和服务器通信
    int fds[2];
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, fds) < 0) {
        std::cerr << "socketpair failed: " << strerror(errno) << std::endl;
        return 1;
    }

    int serverFd = fds[0]; // 服务器端 socket
    int clientFd = fds[1]; // 客户端 socket
    setNonBlocking(serverFd); // 设置非阻塞

    // 准备分散的数据
    const char* header = "HTTP/1.1 200 OK\r\nContent-length: 13\r\n\r\n";
    const char* body = "Hello, World!";
    size_t headerLen = strlen(header); // 38 字节
    size_t bodyLen = strlen(body);    // 13 字节

    // 定义 iovec 数组
    struct iovec iov[2];
    iov[0].iov_base = (void*)header; // 头部数据
    iov[0].iov_len = headerLen;
    iov[1].iov_base = (void*)body;   // 正文数据
    iov[1].iov_len = bodyLen;
    int iovCnt = 2;

    // 使用 writev 写入数据
    ssize_t totalWritten = 0;
    ssize_t toWrite = headerLen + bodyLen; // 总共 51 字节
    while (toWrite > 0) {
        ssize_t len = writev(serverFd, iov, iovCnt);
        if (len < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                std::cout << "Buffer full, retry later" << std::endl;
                break;
            }
            std::cerr << "writev failed: " << strerror(errno) << std::endl;
            close(serverFd);
            close(clientFd);
            return 1;
        }

        totalWritten += len;
        std::cout << "Wrote " << len << " bytes" << std::endl;

        // 更新 iovec
        if (static_cast<size_t>(len) >= iov[0].iov_len) {
            // 第一个块已全部写入，调整第二个块
            size_t excess = len - iov[0].iov_len;
            iov[1].iov_base = (char*)iov[1].iov_base + excess;
            iov[1].iov_len -= excess;
            iov[0].iov_len = 0; // 第一个块已写完
        } else {
            // 只写了一部分第一个块
            iov[0].iov_base = (char*)iov[0].iov_base + len;
            iov[0].iov_len -= len;
        }
        toWrite -= len;

        // 检查是否写完
        if (iov[0].iov_len == 0 && iov[1].iov_len == 0) {
            break;
        }
    }

    std::cout << "Total written: " << totalWritten << " bytes" << std::endl;

    // 从客户端读取数据，验证写入
    char buffer[128] = {0};
    ssize_t readLen = read(clientFd, buffer, sizeof(buffer) - 1);
    if (readLen > 0) {
        std::cout << "Received by client: " << buffer << std::endl;
    } else {
        std::cerr << "read failed: " << strerror(errno) << std::endl;
    }

    // 清理
    close(serverFd);
    close(clientFd);
    return 0;
}